#include "tree.h"
#include "alignment.h"
#include "MbRandom.h"
#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include <istream>
#include <sstream>
#include <fstream>

using namespace std;



Node::Node(void) {

	lft    = NULL;
	rht    = NULL;
	anc    = NULL;
	index  = 0;
	isLeaf = false;
	flag   = false;
	name   = "";
	
}

Tree::Tree(Alignment *al, MbRandom *rp, string treeStr) {

	nodes = NULL;
	downPassSequence = NULL;
	alignmentPtr = al;
	ranPtr = rp;
	
	buildTreeFromNewickString(treeStr);
	
}

Tree::~Tree(void) {

	if (nodes != NULL)
		delete [] nodes;
	if (downPassSequence != NULL)
		delete [] downPassSequence;

}

/* build a tree up from a string with the Newick description of the tree */
void Tree::buildTreeFromNewickString(string treeStr) {

	/* place all of the tree tokens into a string vector, and count the number of taxa in the tree while we are at it */
	vector<string> words;
	numTaxa = 0;
	int numLeftParantheses = 0;
	string nameStr = "";
	bool readingLength = false;
	bool brlensPresent = false;
	for (int i=0; i<treeStr.length(); i++)
		{
		char ch = treeStr.at(i);
		if ( ch == '(' || ch == ')' || ch == ',' || ch == ':' || ch == ';' )
			{
			string tmp = "";
			tmp += ch;
			words.push_back(tmp);
			nameStr = "";
			if ( ch == ':' )
				{
				readingLength = true;
				brlensPresent = true;
				}
			else if ( ch == '(' )
				numLeftParantheses++;
			//cout << "\"" << ch << "\"" << endl;
			}
		else
			{
			nameStr += ch;
			if ( treeStr.at(i+1) == '(' || treeStr.at(i+1) == ')' || treeStr.at(i+1) == ',' || treeStr.at(i+1) == ':' || treeStr.at(i+1) == ';' )
				{
				words.push_back(nameStr);
				if ( readingLength == false )
					numTaxa++;
				else
					readingLength = false;
				}
			}
		}

	/* decide whether the tree is rooted or unrooted from the number of taxa and the number of left parantheses */
	if ( numLeftParantheses == numTaxa - 1 )
		isRooted = true;
	else if ( numLeftParantheses == numTaxa - 2 )
		isRooted = false;
	else
		{
		cerr << "Inconsistent number of parantheses in tree description" << endl;
		exit(1);
		}
		
	/* calculate the number of nodes in the tree */
	if (isRooted == true)
		numNodes = 2 * numTaxa;
	else
		numNodes = 2 * numTaxa - 2;
		
	/* check for consistency in the number of taxa between the data and the tree */
	if ( numTaxa != alignmentPtr->getNumTaxa() )
		{
		cerr << "Problem reading the correct number of taxa." << endl;
		exit(1);
		}

	/* allocate the nodes for the tree */
	nodes = new Node[numNodes];
	for (int i=0; i<numNodes; i++)
		nodes[i].setIndex(i);
	downPassSequence = new Node*[numNodes];
	for (int i=0; i<numNodes; i++)
		downPassSequence[i] = NULL;

	/* build the tree */
	int nextIntNode = numTaxa;
	readingLength = false;
	Node *p, *q;
	for (vector<string>::iterator s=words.begin(); s != words.end(); s++)
		{
		if ( (*s) == "(" )
			{
			/* add an internal node */
			if (nextIntNode == numTaxa)
				{
				p = &nodes[nextIntNode++];
				root = p;
				}
			else
				{
				if (p->getLft() == NULL)
					{
					q = &nodes[nextIntNode++];
					q->setAnc(p);
					p->setLft(q);
					p = q;
					}
				else if (p->getRht() == NULL)
					{
					q = &nodes[nextIntNode++];
					q->setAnc(p);
					p->setRht(q);
					p = q;
					}
				else if (p->getAnc() == NULL)
					{
					q = &nodes[nextIntNode++];
					q->setLft(p);
					p->setAnc(q);
					q->setFlag( true );
					p = q;
					root = p;
					}
				else
					{
					cout << "Problem finding an open node" << endl;
					exit(1);
					}
				}
			}
		else if ( (*s) == ")" )
			{
			/* go down */
			if (p->getFlag() == false)
				p = p->getAnc();
			else
				p = p->getLft();
			}
		else if ( (*s) == "," )
			{
			/* go down */
			if (p->getFlag() == false)
				p = p->getAnc();
			else
				p = p->getLft();
			}
		else if ( (*s) == ":" )
			{
			/* start reading a digit */
			readingLength = true;
			}
		else if ( (*s) == ";" )
			{
			/* end of the tree */

			}
		else
			{
			if ( readingLength == true )
				{
				/* add a branch length */
				double x;
				istringstream buf( (*s) );
				buf >> x;
				//p->setV( x );
				readingLength = false;
				}
			else
				{
				/* add a tip node */
				int taxonIndex = alignmentPtr->getTaxonIndex( (*s) );
				if (taxonIndex == -1)
					{
					cerr << "Couldn't find \"" << (*s) << "\" in the list of taxon names." << endl;
					exit(1);
					}
				
				if (p->getLft() == NULL)
					{
					q = &nodes[taxonIndex];
					q->setAnc(p);
					p->setLft(q);
					p = q;
					}
				else if (p->getRht() == NULL)
					{
					q = &nodes[taxonIndex];
					q->setAnc(p);
					p->setRht(q);
					p = q;
					}
				else if (p->getAnc() == NULL)
					{
					q = &nodes[taxonIndex];
					q->setLft(p);
					p->setAnc(q);
					q->setFlag( true );
					p = q;
					root = p;
					}
				else
					{
					cerr << "Problem finding an open node!" << endl;
					exit(1);
					}
				p->setIsLeaf( true );
				p->setName( *s );
				}
			}
		}
		
	if (isRooted == true)
		{
		p = &nodes[nextIntNode];
		q = root;
		p->setLft( q );
		q->setAnc( p );
		root = p;
		}
		
	/* initialize the post-order tree traversal sequence */
	getDownPassSequence();
	
	for (int n=0; n<numNodes; n++)
		{
		Node *p = getDownPassNode( n );
		if (p->getAnc() != NULL)
			p->setV( ranPtr->exponentialRv(10.0) );
		else
			p->setV( 0.0 );
		}

}

/* get the post-order tree traversal sequence */
void Tree::getDownPassSequence(void) {

	int i = 0;
	passDn(root, &i);

}

/* recursive function filling in the post-order traversal */
void Tree::passDn(Node *p, int *x) {

	if (p != NULL)
		{
		passDn(p->getLft(), x);
		passDn(p->getRht(), x);
		downPassSequence[(*x)++] = p;
		}
		
}

/* print out the tree in an indented format */
void Tree::print(void) {

	showNodes(root, 3);

}

/* show the nodes of the tree in an indented format */
void Tree::showNodes(Node *p, int indent) {

	if (p != NULL)
		{
		for (int i=0; i<indent; i++)
			cout << " ";
		cout << dex(p) << " (" << dex(p->getLft()) << ", " << dex(p->getRht()) << ", " << dex(p->getAnc()) << ") ";
		if ( p != root )
			cout << fixed << setprecision(6) << p->getV();
		if (p->getIsLeaf() == true)
			cout << " (" << p->getName() << ")";
		if (p == root)
			cout << " <- Root" << endl;
		else
			cout << endl;
		showNodes (p->getLft(),  indent + 2);
		showNodes (p->getRht(), indent + 2);
		}
   
}

/* a small function that returns the index of the node (or -1 if the node is NULL) */
int Tree::dex(Node *p) {

	return ( p == NULL ? -1 : p->getIndex() );

}
